package com.ciccwm.log.util;

import lombok.extern.slf4j.Slf4j;
import org.apache.dubbo.rpc.RpcContext;
import org.apache.dubbo.rpc.RpcContextAttachment;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.util.Collections;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.Map;

/**
 * 请求上下文工具类
 * 支持Web请求和Dubbo请求的上下文获取
 */
@Slf4j
public class RequestContextUtil {

    private RequestContextUtil() {
    }

    /**
     * 获取请求头信息
     * 支持Web请求和Dubbo请求
     *
     * @return 请求头信息
     */
    public static Map<String, String> getHeaders() {
        // 尝试获取Web请求上下文
        try {
            HttpServletRequest request = getWebRequest();
            if (request != null) {
                return getWebHeaders(request);
            }

            // 尝试获取Dubbo请求上下文
            return getDubboHeaders();
        } catch (Exception e) {
            log.error("Unexpected error while getting request headers", e);
            return Collections.emptyMap();
        }
    }

    /**
     * 获取Web请求上下文
     */
    private static HttpServletRequest getWebRequest() {
        try {
            ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
            return attributes != null ? attributes.getRequest() : null;
        } catch (Exception e) {
            log.debug("No web request context found");
            return null;
        }
    }

    /**
     * 获取Web请求头信息
     */
    private static Map<String, String> getWebHeaders(HttpServletRequest request) {
        Map<String, String> headers = new HashMap<>();
        Enumeration<String> headerNames = request.getHeaderNames();
        while (headerNames != null && headerNames.hasMoreElements()) {
            String headerName = headerNames.nextElement();
            headers.put(headerName, request.getHeader(headerName));
        }
        return headers;
    }

    /**
     * 获取Dubbo请求头信息
     * 适配Dubbo 3.3.4版本
     */
    private static Map<String, String> getDubboHeaders() {
        try {
            // 获取服务消费者上下文
            RpcContextAttachment consumerContext = RpcContext.getClientAttachment();
            // 获取服务提供者上下文
            RpcContextAttachment providerContext = RpcContext.getServerAttachment();
            
            Map<String, String> headers = new HashMap<>();
            
            // 合并消费者和提供者的上下文信息
            if (consumerContext != null) {
                headers.putAll(consumerContext.getAttachments());
            }
            if (providerContext != null) {
                headers.putAll(providerContext.getAttachments());
            }
            return headers;
        } catch (Exception e) {
            log.debug("No dubbo request context found: {}", e.getMessage());
            return Collections.emptyMap();
        }
    }
} 